library("readr")
library("dplyr")
library("ggplot2")
library("stargazer")
library("patchwork")
library("MASS")

std_error <- function(x) sd(x) / sqrt(length(x))

beer <- read_csv("servingsize.csv")

# Validate numbers
beer |>
    group_by(size) |>
    summarise(consumed_count = mean(consumed / count))

# Get total consumption
sum(beer$consumed)

# Conduct statistical tests
t.test(beer$consumed ~ beer$treatment)

t.test(beer$count ~ beer$treatment)

beer_main <- beer |>
    group_by(size) |>
    summarise(consumed_est = mean(consumed, na.rm = TRUE),
              consumed_ci = std_error(consumed) * 1.96,
              count_est = mean(count, na.rm = TRUE),
              count_ci = std_error(count) * 1.96) |>
    mutate(size = factor(size, levels = c("Standard serving size", "Small serving size")))

beer_main |>
    ggplot(aes(x = size, y = consumed_est,
               ymin = consumed_est - consumed_ci,
               ymax = consumed_est + consumed_ci)) +
    geom_point(size = 2.5) +
    geom_text(aes(label = round(consumed_est, 2)), nudge_x = 0.2) +
    ylim(70, 125) +
    geom_errorbar(width = 0) +
    labs(y = "Average volume of beer (cl)",
         x = NULL) +
    theme_minimal()

ggsave("fig3.png", width = 5, height = 5, bg = "white")

beer_main |>
    ggplot(aes(x = size, y = count_est,
               ymin = count_est - count_ci,
               ymax = count_est + count_ci)) +
    geom_point(size = 2.5) +
    geom_text(aes(label = round(count_est, 2)), nudge_x = 0.2) +
    ylim(1.5, 3) +
    geom_errorbar(width = 0) +
    labs(y = "Average number of servings",
         x = NULL) +
    theme_minimal()

ggsave("fig4.png", width = 5, height = 5, bg = "white")

beer_gender <- beer |>
    group_by(size, male) |>
    summarise(consumed_est = mean(consumed, na.rm = TRUE),
              consumed_ci = std_error(consumed) * 1.96,
              count_est = mean(count, na.rm = TRUE),
              count_ci = std_error(count) * 1.96,
              .groups = "drop") |>
    mutate(male = ifelse(male == 1, "Male", "Female"),
           size = factor(size, levels = c("Standard serving size", "Small serving size")))

fig_gender_consumed <- beer_gender |>
    ggplot(aes(x = size, y = consumed_est,
               ymin = consumed_est - consumed_ci,
               ymax = consumed_est + consumed_ci)) +
    geom_point() +
    geom_errorbar(width = 0) +
    geom_text(aes(label = round(consumed_est, 2)), nudge_x = 0.2) +
    facet_wrap(~ male) +
    labs(y = "Average volume of beer (cl)",
         x = NULL) +
    theme_bw()

fig_gender_count <- beer_gender |>
    ggplot(aes(x = size, y = count_est,
               ymin = count_est - count_ci,
               ymax = count_est + count_ci)) +
    geom_point() +
    geom_errorbar(width = 0) +
    geom_text(aes(label = round(count_est, 2)), nudge_x = 0.2) +
    facet_wrap(~ male) +
    labs(y = "Average number of servings",
         x = NULL) +
    theme_bw()

fig_gender_consumed / fig_gender_count

ggsave("fig5.png", width = 8, height = 6, bg = "white")

reg_1 <- lm(consumed ~ treatment, data = beer)
reg_2 <- lm(count ~ treatment, data = beer)
reg_3 <- lm(consumed ~ treatment * male, data = beer)
reg_4 <- lm(count ~ treatment * male, data = beer)

summary(glm.nb(count ~ treatment, data = beer))

stargazer(reg_1, reg_2, reg_3, reg_4,
          column.labels = c("Volume, average", "Servings, average", "Volume, interaction", "Servings, interaction"),
          covariate.labels = c("Small serving size", "Male", "Small serving size * Male"),
          digits = 2,
          align = TRUE,
          no.space = TRUE,
          model.names = FALSE,
          notes.align = "l",
          keep.stat = c("n", "rsq"),
          out = "reg_table.htm",
          type = "text")

sink("sessionInfo.txt")
cat("\nThe results are produced with 01-analysis.R with this session in R:\n\n")
sessionInfo()
sink()
